package com.ruoyi.demo.config;

import com.baomidou.dynamic.datasource.provider.DynamicDataSourceProvider;
import com.google.common.base.Preconditions;
import com.ruoyi.common.core.exception.ServiceException;
import lombok.RequiredArgsConstructor;

import org.apache.shardingsphere.api.config.sharding.ShardingRuleConfiguration;
import org.apache.shardingsphere.api.config.sharding.TableRuleConfiguration;
import org.apache.shardingsphere.core.yaml.swapper.ShardingRuleConfigurationYamlSwapper;
import org.apache.shardingsphere.shardingjdbc.api.ShardingDataSourceFactory;
import org.apache.shardingsphere.shardingjdbc.spring.boot.common.SpringBootPropertiesConfigurationProperties;
import org.apache.shardingsphere.shardingjdbc.spring.boot.encrypt.SpringBootEncryptRuleConfigurationProperties;
import org.apache.shardingsphere.shardingjdbc.spring.boot.masterslave.SpringBootMasterSlaveRuleConfigurationProperties;
import org.apache.shardingsphere.shardingjdbc.spring.boot.sharding.SpringBootShardingRuleConfigurationProperties;
import org.apache.shardingsphere.spring.boot.datasource.DataSourcePropertiesSetter;
import org.apache.shardingsphere.spring.boot.datasource.DataSourcePropertiesSetterHolder;
import org.apache.shardingsphere.spring.boot.util.DataSourceUtil;
import org.apache.shardingsphere.spring.boot.util.PropertyUtil;
import org.apache.shardingsphere.underlying.common.config.inline.InlineExpressionParser;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.ComponentScan;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;
import org.springframework.core.env.StandardEnvironment;
import org.springframework.jndi.JndiObjectFactoryBean;

import javax.naming.NamingException;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.*;

//@Configuration
@ComponentScan("org.apache.shardingsphere.spring.boot.converter")
@EnableConfigurationProperties({
    SpringBootShardingRuleConfigurationProperties.class,
    SpringBootMasterSlaveRuleConfigurationProperties.class, SpringBootEncryptRuleConfigurationProperties.class,
    SpringBootPropertiesConfigurationProperties.class})
@RequiredArgsConstructor
public class ShardingConfiguration implements EnvironmentAware {

    private final SpringBootShardingRuleConfigurationProperties shardingRule;

    private final SpringBootPropertiesConfigurationProperties props;

    private final String jndiName = "jndi-name";

    private final Map<String, DataSource> dataSourceMap = new LinkedHashMap<>();


    @Bean
    public ShardingDataSourceMap shardingDataSourceMap() throws SQLException {
        ShardingRuleConfiguration shardingRuleConfiguration = new ShardingRuleConfigurationYamlSwapper().swap(shardingRule);
        Map<String, DataSource> shardingDataSourceMap = new LinkedHashMap<>();
        Collection<TableRuleConfiguration> tableRuleConfigs = shardingRuleConfiguration.getTableRuleConfigs();
        for (String dataSourceKey : dataSourceMap.keySet()) {
            Map<String, DataSource> map = new HashMap<>();
            map.put(dataSourceKey, dataSourceMap.get(dataSourceKey));
            List<TableRuleConfiguration> newShardingRuleConfigurationList = new LinkedList<>();
            for (TableRuleConfiguration tableRuleConfig : tableRuleConfigs) {
                TableRuleConfiguration tableRuleConfiguration = new TableRuleConfiguration(tableRuleConfig.getLogicTable(), dataSourceKey + "." + tableRuleConfig.getActualDataNodes());
                tableRuleConfiguration.setDatabaseShardingStrategyConfig(tableRuleConfig.getDatabaseShardingStrategyConfig());
                tableRuleConfiguration.setTableShardingStrategyConfig(tableRuleConfig.getTableShardingStrategyConfig());
                tableRuleConfiguration.setKeyGeneratorConfig(tableRuleConfig.getKeyGeneratorConfig());
                newShardingRuleConfigurationList.add(tableRuleConfiguration);
            }
            shardingRuleConfiguration.setTableRuleConfigs(newShardingRuleConfigurationList);
            DataSource shardingDataSource = ShardingDataSourceFactory.createDataSource(map,
                shardingRuleConfiguration,
                props.getProps());
            shardingDataSourceMap.put(dataSourceKey, shardingDataSource);
        }
        return new ShardingDataSourceMap(shardingDataSourceMap);
    }

    @Bean
    public DynamicDataSourceProvider dynamicDataSourceProvider(ShardingDataSourceMap shardingDataSourceMap) {
        return new DynamicDataSourceProvider() {

            @Override
            public Map<String, DataSource> loadDataSources() {

                return shardingDataSourceMap.getDataSourceMap();
            }
        };
    }

    @Override
    public final void setEnvironment(final Environment environment) {
        String prefix = "spring.shardingsphere.datasource.";
        for (String each : getDataSourceNames(environment, prefix)) {
            try {
                dataSourceMap.put(each, getDataSource(environment, prefix, each));
            } catch (final ReflectiveOperationException ex) {
                throw new ServiceException("Can't find datasource type!");
            } catch (final NamingException namingEx) {
                throw new ServiceException("Can't find JNDI datasource!");
            }
        }
    }

    private List<String> getDataSourceNames(final Environment environment, final String prefix) {
        StandardEnvironment standardEnv = (StandardEnvironment) environment;
        standardEnv.setIgnoreUnresolvableNestedPlaceholders(true);
        return null == standardEnv.getProperty(prefix + "name")
            ? new InlineExpressionParser(standardEnv.getProperty(prefix + "names")).splitAndEvaluate() : Collections.singletonList(standardEnv.getProperty(prefix + "name"));
    }

    @SuppressWarnings("unchecked")
    private DataSource getDataSource(final Environment environment, final String prefix, final String dataSourceName) throws ReflectiveOperationException, NamingException {
        Map<String, Object> dataSourceProps = PropertyUtil.handle(environment, prefix + dataSourceName.trim(), Map.class);
        Preconditions.checkState(!dataSourceProps.isEmpty(), "Wrong datasource properties!");
        if (dataSourceProps.containsKey(jndiName)) {
            return getJndiDataSource(dataSourceProps.get(jndiName).toString());
        }
        DataSource result = DataSourceUtil.getDataSource(dataSourceProps.get("type").toString(), dataSourceProps);
        Optional<DataSourcePropertiesSetter> dataSourcePropertiesSetter = DataSourcePropertiesSetterHolder.getDataSourcePropertiesSetterByType(dataSourceProps.get("type").toString());
        if (dataSourcePropertiesSetter.isPresent()) {
            dataSourcePropertiesSetter.get().propertiesSet(environment, prefix, dataSourceName, result);
        }
        return result;


    }

    private DataSource getJndiDataSource(final String jndiName) throws NamingException {
        JndiObjectFactoryBean bean = new JndiObjectFactoryBean();
        bean.setResourceRef(true);
        bean.setJndiName(jndiName);
        bean.setProxyInterface(DataSource.class);
        bean.afterPropertiesSet();
        return (DataSource) bean.getObject();
    }
}
